import os
import argparse
from hashlib import pbkdf2_hmac

from Cryptodome.Cipher import AES
from Cryptodome.Random import get_random_bytes


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-f",
        "--file",
        type=str,
        required=True,
        help="Input file.",
    )
    parser.add_argument(
        "-o",
        "--out",
        type=str,
        required=True,
        help="Output file.",
    )
    parser.add_argument(
        "-p",
        "--password",
        type=str,
        required=True,
        help="Password.",
    )
    return parser.parse_args()


def main():
    args = parse_args()

    if not os.path.exists(args.file):
        print("* Invalid filepath")
        return

    with open(args.file, "rb") as file:
        data = file.read()

    salt = get_random_bytes(16)
    dervied_key = pbkdf2_hmac(
        "sha256",
        args.password.encode("utf-8"),
        salt,
        500_000,
    )

    cipher = AES.new(dervied_key, AES.MODE_EAX)
    ciphertext, tag = cipher.encrypt_and_digest(data)

    file_out = open(args.out, "wb")

    list(map(file_out.write, (salt, cipher.nonce, tag, ciphertext)))
    file_out.close()


if __name__ == "__main__":
    main()
